import torch
import torch.nn.functional as F
from network import MLP


def main():
    net = MLP()
    net.load_state_dict(torch.load('../../models/MLP.pth'))

    torch.manual_seed(2)
    n_data = torch.ones(1, 2)
    x0 = torch.normal(1 * n_data, 1)
    # y0 = torch.zeros(1)
   
    net.eval()
    with torch.no_grad():
        outputs = net(x0)
        print(F.softmax(outputs, dim=1))
        predict = torch.max(outputs, dim=1)[1].numpy()
        print(predict)


if __name__ == '__main__':
    main()
